package com.xkcyy.one.ai.agents;

import com.xkcyy.one.ai.common.ToolAgent;
import com.xkcyy.one.ai.common.tool.AbstractSimpleTool;
import com.xkcyy.one.ai.common.tool.AbstractTool;
import com.xkcyy.one.ai.common.tool.Tool;
import com.xkcyy.one.ai.common.tool.ToolArguments;
import com.xkcyy.one.ai.llm.OpenAiClient;
import com.xkcyy.one.ai.llm.model.ChatMessage;
import com.xkcyy.one.ai.llm.model.ChatResult;
import com.xkcyy.one.ai.llm.model.ToolCallResult;
import com.xkcyy.one.ai.llm.model.ToolDefinition;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

@Slf4j
public class CallableToolAgent extends AbstractSimpleTool<CallableToolAgent.CallableToolAgentArguments>
        implements ToolAgent {

    private final OpenAiClient openAiClient;
    private final Map<String, Tool> tools;

    public CallableToolAgent() {
        this.openAiClient = new OpenAiClient();
        this.tools = new ConcurrentHashMap<>();
    }

    public CallableToolAgent(OpenAiClient openAiClient) {
        this.openAiClient = openAiClient;
        this.tools = new ConcurrentHashMap<>();
    }

    @Override
    public void registerTool(Tool toolExecutor) {
        tools.put(toolExecutor.getName(), toolExecutor);
    }

    @Override
    public void registerTools(List<Tool> executors) {
        executors.forEach(this::registerTool);
    }

    @Override
    public Flux<ChatResult> chatStream(List<ChatMessage> messages) {
        // 准备工具列表
        List<ToolDefinition> toolDefs = prepareTools();
        if (toolDefs.isEmpty()) {
            return openAiClient.chatStream(messages);
        }

        // 用于收集工具调用和最终内容
        AtomicReference<String> finalContent = new AtomicReference<>("");
        List<ToolCallResult> collectedToolCalls = new ArrayList<>();
        List<String> collectedToolResults = new ArrayList<>();

        // 使用带工具的流式聊天
        return openAiClient.chatStreamWithTools(messages, toolDefs)
                .doOnNext(response -> {
                    // 累积内容
                    if (response.getContent() != null) {
                        finalContent.updateAndGet(current -> current + response.getContent());
                    }

                    // 收集工具调用
                    if (response.getToolCalls() != null && !response.getToolCalls().isEmpty()) {
                        collectedToolCalls.addAll(response.getToolCalls());
                    }
                })
                .concatMap(response -> {
                    // 如果流还未结束，返回中间结果
                    if (!response.isComplete()) {
                        return Mono.just(response);
                    }

                    // 流结束时，如果有工具调用，执行它们
                    if (!collectedToolCalls.isEmpty()) {
                        List<ChatMessage> updatedMessages = new ArrayList<>(messages);

                        // 添加助手消息
                        ChatMessage assistantMessage = ChatMessage.builder()
                                .role("assistant")
                                .content(finalContent.get())
                                .toolCalls(collectedToolCalls)
                                .build();
                        updatedMessages.add(assistantMessage);

                        // 首先处理所有工具调用，然后继续对话
                        return Flux.fromIterable(collectedToolCalls)
                                .index()
                                .flatMap(tuple -> {
                                    ToolCallResult toolCall = tuple.getT2();
                                    String toolName = toolCall.getFunction().getName();
                                    Tool executor = this.tools.get(toolName);

                                    if (executor == null) {
                                        log.error("找不到工具执行器: {}", toolName);
                                        return Mono.empty();
                                    }

                                    if (executor instanceof ToolAgent) {
                                        Flux<ChatResult> chatResultFlux = ((ToolAgent) executor)
                                                .chatStream(toolCall.getFunction().getArguments());
                                        return chatResultFlux
                                                .map(chatResult -> {
                                                    // 实时转发ToolAgent的结果
                                                    String toolResult = chatResult.getContent();

                                                    // 如果是最后一个结果，添加到消息历史
                                                    if (chatResult.isComplete()) {
                                                        collectedToolResults.add(toolResult);
                                                        ChatMessage toolResponseMessage = ChatMessage.builder()
                                                                .role("tool")
                                                                .content(toolResult)
                                                                .build();
                                                        updatedMessages.add(toolResponseMessage);
                                                    }

                                                    // 返回工具调用结果
                                                    return ChatResult.builder()
                                                            .content("")
                                                            .innerResult(chatResult)
                                                            .toolCalls(Collections.singletonList(toolCall))
                                                            .toolResults(Collections.singletonList(toolResult))
                                                            .isComplete(chatResult.isComplete())
                                                            .build();
                                                });
                                    }

                                    // 执行普通工具
                                    String toolResult = executor.execute(toolCall.getFunction().getArguments());
                                    collectedToolResults.add(toolResult);

                                    // 添加工具响应消息
                                    ChatMessage toolResponseMessage = ChatMessage.builder()
                                            .role("tool")
                                            .content(toolResult)
                                            .build();
                                    updatedMessages.add(toolResponseMessage);

                                    // 返回工具调用结果
                                    return Mono.just(ChatResult.builder()
                                            .content(toolResult)
                                            .toolCalls(Collections.singletonList(toolCall))
                                            .toolResults(Collections.singletonList(toolResult))
                                            .isComplete(false)
                                            .build());
                                })
                                .collectList()
                                .flatMapMany(toolResults -> {
                                    // 所有工具调用处理完成后，继续对话
                                    return this.chatStream(updatedMessages);
                                });
                    }

                    // 如果没有工具调用，返回最终结果
                    return Mono.just(ChatResult.builder()
                            .content(finalContent.get())
                            .reasoningContent(response.getReasoningContent())
                            .toolCalls(collectedToolCalls)
                            .toolResults(collectedToolResults)
                            .usage(response.getUsage())
                            .isComplete(true)
                            .build());
                });
    }

    /**
     * 准备可用的工具列表
     *
     * @return 工具列表
     */
    private List<ToolDefinition> prepareTools() {
        if (tools.isEmpty()) {
            return Collections.emptyList();
        }

        return tools.values().stream()
                .map(Tool::getDefinition)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    /**
     * 处理非流式对话中的工具调用
     *
     * @param messages 原始消息列表
     * @param response 包含工具调用的响应
     * @return 最终处理后的结果
     */
    private ChatResult handleToolCalls(List<ChatMessage> messages, ChatResult response) {
        List<ChatMessage> updatedMessages = new ArrayList<>(messages);
        List<String> toolResults = new ArrayList<>();

        // 添加助手消息
        ChatMessage assistantMessage = ChatMessage.builder()
                .role("assistant")
                .content(response.getContent())
                .build();
        updatedMessages.add(assistantMessage);

        // 依次处理每个工具调用并添加到消息列表
        for (ToolCallResult toolCall : response.getToolCalls()) {
            String toolName = toolCall.getFunction().getName();
            Tool executor = tools.get(toolName);

            if (executor == null) {
                log.error("找不到工具执行器: {}", toolName);
                continue;
            }

            // 执行工具
            String toolResult = executor.execute(toolCall.getFunction().getArguments());
            toolResults.add(toolResult);

            // 添加工具响应消息
            ChatMessage toolResponseMessage = ChatMessage.builder()
                    .role("tool")
                    .content(toolResult)
                    .build();
            updatedMessages.add(toolResponseMessage);
        }

        // 继续对话
        ChatResult result = openAiClient.chat(updatedMessages);
        result.setToolResults(toolResults);
        return result;
    }

    /**
     * 处理流式对话中的工具调用
     *
     * @param messages 原始消息列表
     * @param response 包含工具调用的响应
     * @return 工具执行后继续的流
     */
    private Flux<ChatResult> handleToolCallsStream(List<ChatMessage> messages, ChatResult response) {
        List<ChatMessage> updatedMessages = new ArrayList<>(messages);

        // 添加助手消息
        ChatMessage assistantMessage = ChatMessage.builder()
                .role("assistant")
                .content(response.getContent())
                .build();
        updatedMessages.add(assistantMessage);

        // 处理多个工具调用
        List<ChatMessage> toolResponses = new ArrayList<>();
        for (ToolCallResult toolCall : response.getToolCalls()) {
            String toolName = toolCall.getFunction().getName();
            Tool executor = tools.get(toolName);

            if (executor == null) {
                log.error("找不到工具执行器: {}", toolName);
                continue;
            }

            // 执行工具
            String toolResult = executor.execute(toolCall.getFunction().getArguments());

            // 创建工具响应消息
            ChatMessage toolResponseMessage = ChatMessage.builder()
                    .role("tool")
                    .content(toolResult)
                    .build();
            toolResponses.add(toolResponseMessage);
        }

        // 添加所有工具响应
        updatedMessages.addAll(toolResponses);

        // 继续流式对话
        return this.chatStream(updatedMessages);
    }

    @Override
    public String doExecute(CallableToolAgentArguments arguments) {
        return "";
    }

    @EqualsAndHashCode(callSuper = true)
    @Data
    public static class CallableToolAgentArguments extends ToolArguments {
        private String prompt;
    }
}